#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient
import fire


# 连接MongoDB
def InitMongoDB(MONGODB_HOST, MONGODB_PORT, MONGODB_DB_NAME, MONGODB_COLLECTION_NAME):
    client = MongoClient(MONGODB_HOST, MONGODB_PORT)
    db = client[MONGODB_DB_NAME]
    collection = db[MONGODB_COLLECTION_NAME]
    return client, collection


# 关闭MongoDB
def CloseMongoDB(MongoDB_Connection):
    MongoDB_Connection.close()


def GenProvinceFST(filepath: str):
    client, collection = InitMongoDB("localhost", 27017, "hanvcf", "fst")

    province2subgroup = {}
    province2subgroup["hei_long_jiang"] = "Northeast Han"
    province2subgroup["ji_lin"] = "Northeast Han"
    province2subgroup["liao_ning"] = "Northeast Han"
    province2subgroup["he_bei"] = "Central China Han"
    province2subgroup["bei_jing"] = "Central China Han"
    province2subgroup["tian_jin"] = "Central China Han"
    province2subgroup["shan_dong"] = "Central China Han"
    province2subgroup["he_nan"] = "Central China Han"
    province2subgroup["shan_xi"] = "Central China Han"
    province2subgroup["shaan_xi"] = "Northwest Han"
    province2subgroup["nei_meng_gu"] = "Northwest Han"
    province2subgroup["ning_xia"] = "Northwest Han"
    province2subgroup["gan_su"] = "Northwest Han"
    province2subgroup["qing_hai"] = "Northwest Han"
    province2subgroup["xin_jiang"] = "Northwest Han"
    province2subgroup["jiang_su"] = "Southeast Han"
    province2subgroup["an_hui"] = "Southeast Han"
    province2subgroup["shang_hai"] = "Southeast Han"
    province2subgroup["zhe_jiang"] = "Southeast Han"
    province2subgroup["hu_bei"] = "Southwest Han"
    province2subgroup["jiang_xi"] = "Southeast Han"
    province2subgroup["hu_nan"] = "Southwest Han"
    province2subgroup["chong_qing"] = "Southwest Han"
    province2subgroup["si_chuan"] = "Southwest Han"
    province2subgroup["gui_zhou"] = "Southwest Han"
    province2subgroup["yun_nan"] = "Southwest Han"
    province2subgroup["xi_zang"] = "Southwest Han"
    province2subgroup["fu_jian"] = "South coast Han"
    province2subgroup["guang_dong"] = "South coast Han"
    province2subgroup["guang_xi"] = "South coast Han"
    province2subgroup["tai_wan"] = "South coast Han"
    province2subgroup["hai_nan"] = "South coast Han"
    province2subgroup["xiang_gang"] = "South coast Han"
    province2subgroup["singapore"] = "South coast Han"

    type = "province"
    infile = open(filepath, "r")
    title = infile.readline().rstrip().split("\t")
    for line in infile:
        line = line.rstrip()
        splitline = line.split("\t")
        name = splitline[0].strip()
        genetic_subgroup = province2subgroup[name]
        fst_data = []
        for i in range(1, len(splitline)):
            name_data = title[i].strip()
            genetic_subgroup_data = province2subgroup[name_data]
            value_data = float(splitline[i])
            fst_data.append({"name": name_data, "value": value_data, "genetic_subgroup": genetic_subgroup_data})
        collection.insert_one({"name": name, "type": type, "genetic_subgroup": genetic_subgroup, "fst_data": fst_data})

    CloseMongoDB(client)
    infile.close()


def GenGeneticSubgroupFST(filepath: str):
    client, collection = InitMongoDB("localhost", 27017, "hanvcf", "fst")
    type = "genetic_subgroup"
    infile = open(filepath, "r")
    title = infile.readline().rstrip().split("\t")

    gs = set(
        ["Northeast Han", "Central China Han", "Northwest Han", "Southeast Han", "Southwest Han", "South coast Han"])

    for line in infile:
        line = line.rstrip()
        splitline = line.split("\t")
        name = splitline[0].strip()
        fst_data = []
        for i in range(1, len(splitline)):
            name_data = title[i].strip()
            type_data = "genetic_subgroup" if (title[i] in gs) else "Population"
            value_data = float(splitline[i])
            fst_data.append({"name": name_data, "value": value_data, "type": type_data})
        collection.insert_one({"name": name, "type": type, "fst_data": fst_data})
    CloseMongoDB(client)
    infile.close()


if __name__ == '__main__':
    fire.Fire()
